# 导入必要的库
import os  # 操作系统接口模块
import pandas as pd  # 数据处理库
import numpy as np  # 数值计算库
import matplotlib.pyplot as plt  # 绘图库
import seaborn as sns  # 基于matplotlib的统计数据可视化库
from typing import Dict, List, Tuple, Optional  # 类型提示支持
import argparse  # 命令行参数解析

# 从本地模块导入函数
from data_loader import load_data, extract_class_from_filename


def plot_sensor_data(df: pd.DataFrame, title: str = 'Sensor Data', 
                    output_path: Optional[str] = None):
    """
    绘制传感器数据图表
    
    Args:
        df: 包含传感器数据的DataFrame
        title: 图表标题
        output_path: 图表保存路径，如果为None则显示图表
    """
    # 创建14x10英寸的图表
    plt.figure(figsize=(14, 10))
    
    # 绘制加速度数据子图
    plt.subplot(2, 1, 1)  # 2行1列的第1个子图
    plt.plot(df['acc_x'], label='acc_x')  # 绘制x轴加速度
    plt.plot(df['acc_y'], label='acc_y')  # 绘制y轴加速度
    plt.plot(df['acc_z'], label='acc_z')  # 绘制z轴加速度
    plt.title(f'{title} - Acceleration')  # 设置子图标题
    plt.xlabel('Sample')  # 设置x轴标签
    plt.ylabel('Acceleration')  # 设置y轴标签
    plt.legend()  # 显示图例
    plt.grid(True)  # 显示网格
    
    # 绘制角度数据子图
    plt.subplot(2, 1, 2)  # 2行1列的第2个子图
    plt.plot(df['angle_x'], label='angle_x')  # 绘制x轴角度
    plt.plot(df['angle_y'], label='angle_y')  # 绘制y轴角度
    plt.plot(df['angle_z'], label='angle_z')  # 绘制z轴角度
    plt.title(f'{title} - Angle')  # 设置子图标题
    plt.xlabel('Sample')  # 设置x轴标签
    plt.ylabel('Angle (degrees)')  # 设置y轴标签
    plt.legend()  # 显示图例
    plt.grid(True)  # 显示网格
    
    plt.tight_layout()  # 自动调整子图参数
    
    # 根据output_path决定保存还是显示图表
    if output_path:
        plt.savefig(output_path)  # 保存图表到文件
        plt.close()  # 关闭图表
    else:
        plt.show()  # 显示图表


def plot_class_comparison(data_by_class: Dict[str, pd.DataFrame], 
                         feature: str = 'acc_z', output_path: Optional[str] = None):
    """
    绘制不同类别间某个特征的比较图表
    
    Args:
        data_by_class: 类别名到DataFrame的映射字典
        feature: 要比较的特征
        output_path: 图表保存路径，如果为None则显示图表
    """
    plt.figure(figsize=(14, 8))  # 创建14x8英寸的图表
    
    for class_name, df in data_by_class.items():
        # 对数据进行降采样以使图表更清晰
        sample_rate = max(1, len(df) // 1000)  # 计算采样率
        df_sampled = df.iloc[::sample_rate]  # 按采样率降采样
        
        plt.plot(range(len(df_sampled)), df_sampled[feature], label=class_name)  # 绘制特征数据
    
    plt.title(f'Comparison of {feature} across classes')  # 设置图表标题
    plt.xlabel('Sample (downsampled)')  # 设置x轴标签
    plt.ylabel(feature)  # 设置y轴标签
    plt.legend()  # 显示图例
    plt.grid(True)  # 显示网格
    
    if output_path:
        plt.savefig(output_path)  # 保存图表到文件
        plt.close()  # 关闭图表
    else:
        plt.show()  # 显示图表


def plot_feature_distributions(data_by_class: Dict[str, pd.DataFrame], 
                              features: List[str], output_dir: Optional[str] = None):
    """
    绘制不同类别间特征的分布图表
    
    Args:
        data_by_class: 类别名到DataFrame的映射字典
        features: 要绘制的特征列表
        output_dir: 图表保存目录，如果为None则显示图表
    """
    if output_dir:
        os.makedirs(output_dir, exist_ok=True)  # 创建输出目录
    
    for feature in features:
        plt.figure(figsize=(12, 6))  # 创建12x6英寸的图表
        
        for class_name, df in data_by_class.items():
            sns.kdeplot(df[feature], label=class_name)  # 绘制核密度估计图
        
        plt.title(f'Distribution of {feature} across classes')  # 设置图表标题
        plt.xlabel(feature)  # 设置x轴标签
        plt.ylabel('Density')  # 设置y轴标签
        plt.legend()  # 显示图例
        plt.grid(True)  # 显示网格
        
        if output_dir:
            output_path = os.path.join(output_dir, f'{feature}_distribution.png')  # 构建输出路径
            plt.savefig(output_path)  # 保存图表到文件
            plt.close()  # 关闭图表
        else:
            plt.show()  # 显示图表


def create_heatmap(data_by_class: Dict[str, pd.DataFrame], 
                  output_path: Optional[str] = None):
    """
    创建特征相关性热力图
    
    Args:
        data_by_class: 类别名到DataFrame的映射字典
        output_path: 图表保存路径，如果为None则显示图表
    """
    # 合并所有类别的数据
    all_data = pd.concat(list(data_by_class.values()))
    
    # 选择传感器相关列
    sensor_cols = ['acc_x', 'acc_y', 'acc_z', 'angle_x', 'angle_y', 'angle_z']
    
    # 计算相关系数矩阵
    corr_matrix = all_data[sensor_cols].corr()
    
    plt.figure(figsize=(10, 8))  # 创建10x8英寸的图表
    sns.heatmap(corr_matrix, annot=True, cmap='coolwarm', vmin=-1, vmax=1)  # 绘制热力图
    plt.title('Feature Correlation Heatmap')  # 设置图表标题
    
    if output_path:
        plt.savefig(output_path)  # 保存图表到文件
        plt.close()  # 关闭图表
    else:
        plt.show()  # 显示图表


def visualize_data(data_dir: str, output_dir: Optional[str] = None):
    """
    可视化传感器数据
    
    Args:
        data_dir: 包含数据文件的目录路径
        output_dir: 图表保存目录，如果为None则显示图表
    """
    # 如果指定了输出目录则创建
    if output_dir:
        os.makedirs(output_dir, exist_ok=True)
    
    # 加载数据
    data_by_class, classes = load_data(data_dir)
    print(f"Loaded data for {len(classes)} classes: {classes}")
    
    # 为每个类别绘制样本数据
    for class_name, df in data_by_class.items():
        output_path = None
        if output_dir:
            output_path = os.path.join(output_dir, f'{class_name}_sensor_data.png')
        
        plot_sensor_data(df.iloc[:500], title=f'Class: {class_name}', output_path=output_path)
    
    # 绘制不同类别间特征的比较
    for feature in ['acc_x', 'acc_y', 'acc_z', 'angle_x', 'angle_y', 'angle_z']:
        output_path = None
        if output_dir:
            output_path = os.path.join(output_dir, f'{feature}_comparison.png')
        
        plot_class_comparison(data_by_class, feature=feature, output_path=output_path)
    
    # 绘制特征分布
    if output_dir:
        plot_feature_distributions(
            data_by_class, 
            ['acc_x', 'acc_y', 'acc_z', 'angle_x', 'angle_y', 'angle_z'],
            output_dir
        )
    
    # 创建热力图
    output_path = None
    if output_dir:
        output_path = os.path.join(output_dir, 'correlation_heatmap.png')
    
    create_heatmap(data_by_class, output_path=output_path)


if __name__ == "__main__":
    # 设置命令行参数解析
    parser = argparse.ArgumentParser(description='Visualize sensor data')
    parser.add_argument('--data_dir', type=str, default='../data',
                        help='Path to directory containing data files')
    parser.add_argument('--output_dir', type=str, default='../visualizations',
                        help='Directory to save the plots')
    
    args = parser.parse_args()  # 解析命令行参数
    
    visualize_data(args.data_dir, args.output_dir)  # 调用可视化函数
